import pandas as pd
import numpy as np 
from typing import List, Literal
import matplotlib.pyplot as plt # 可视化

from sklearn.tree import DecisionTreeClassifier as dtc # 树算法
from sklearn.model_selection import train_test_split # 拆分数据
from sklearn.metrics import accuracy_score # 模型准确度
from sklearn.tree import plot_tree # 树图

def decision_tree(df: pd.DataFrame, feature_names:List[str], class_name:str, criterion:Literal['gini', 'entropy', 'log_loss'] ='gini') -> dtc:
    '''
    传入数据的 dataframe,  特征的标签和分类类别
    '''
    # 拆分数据集
    X_train, X_test, y_train, y_test = train_test_split(df[feature_names], df[class_name], test_size=0.2, random_state=42)

    # 训练模型
    # 使用 entropy 或者 gini
    model = dtc(criterion=criterion, max_depth=5, random_state=42)
    model.fit(X_train, y_train)

    # 预测
    y_pred = model.predict(X_test)

    # 评估
    accuracy = accuracy_score(y_test, y_pred)
    print(f"决策树模型准确度: {accuracy * 100:.1f}%")

    # 可视化
    plt.figure(figsize=(12, 8))
    plot_tree(model, feature_names=feature_names,class_names=[str(c) for c in model.classes_], filled=True)
    plt.show()

    return model

    
if __name__ == "__main__":
    # 创建模拟药物数据
    np.random.seed(42)
    n_samples = 200

    # 生成模拟数据
    data = {
        'Age': np.random.randint(20, 80, n_samples),
        'Sex': np.random.choice(['M', 'F'], n_samples),
        'BP': np.random.choice(['LOW', 'NORMAL', 'HIGH'], n_samples),
        'Cholesterol': np.random.choice(['NORMAL', 'HIGH'], n_samples),
        'Na_to_K': np.random.uniform(6, 40, n_samples)
    }

    # 创建药物标签（基于一些规则来模拟真实情况）
    drugs = []
    for i in range(n_samples):
        age = data['Age'][i]
        bp = data['BP'][i]
        chol = data['Cholesterol'][i]
        na_k = data['Na_to_K'][i]
        
        if bp == 'HIGH' and chol == 'HIGH':
            drug = 'drugY'
        elif bp == 'HIGH' and age > 50:
            drug = 'drugC'
        elif bp == 'NORMAL' and na_k > 25:
            drug = 'drugX'
        elif bp == 'LOW':
            drug = 'drugA'
        else:
            drug = 'drugB'
        
        drugs.append(drug)

    data['Drug'] = drugs

    # 创建DataFrame
    df = pd.DataFrame(data)

    print("模拟数据集信息:")
    print(df.head(10))
    print(f"\n数据形状: {df.shape}")
    print(f"\n药物分布:")
    print(df['Drug'].value_counts())

    # 数据预处理 - 将分类变量转换为数值
    df_processed = df.copy()

    # 性别编码
    df_processed['Sex'] = df_processed['Sex'].map({'M': 0, 'F': 1})

    # 血压编码
    df_processed['BP'] = df_processed['BP'].map({'LOW': 0, 'NORMAL': 1, 'HIGH': 2})

    # 胆固醇编码
    df_processed['Cholesterol'] = df_processed['Cholesterol'].map({'NORMAL': 0, 'HIGH': 1})

    print(f"\n预处理后的数据:")
    print(df_processed.head())

    # 准备特征和目标变量
    feature_names = ['Age', 'Sex', 'BP', 'Cholesterol', 'Na_to_K']
    class_names = df_processed['Drug'].unique().astype(str).tolist()

    model = decision_tree(df=df_processed, feature_names=feature_names, class_name='Drug', criterion='entropy')

    # 打印特征重要性
    feature_importance = model.feature_importances_
    print(f"\n特征重要性:")
    for i, importance in enumerate(feature_importance):
        print(f"{feature_names[i]}: {importance:.3f}")

    # 示例预测
    print(f"\n预测示例:")
    sample_cases = [
        [45, 0, 2, 1, 30.5],  # 45岁男性，高血压，高胆固醇 -> drugY
        [25, 1, 0, 0, 15.2],  # 25岁女性，低血压，正常胆固醇 -> drugA
        [60, 0, 1, 0, 20.1],  # 60岁男性，正常血压，正常胆固醇 -> drugB
        [60, 0, 2, 0, 35.0],   # 60岁男性，高血压，低固醇 -> drugC
        [60, 0, 1, 0, 26.1],  # 60岁男性，正常血压，正常胆固醇 -> drugX
    ]

    sample_df = pd.DataFrame(sample_cases, columns=feature_names)

    # 批量预测
    predictions = model.predict(sample_df)
    probabilities = model.predict_proba(sample_df)

    for i, (case, pred, prob) in enumerate(zip(sample_cases, predictions, probabilities)):
        max_prob = max(prob)
        print(f"案例 {i+1}: {case} -> 预测药物: {pred} (概率: {max_prob:.2f})")